"""
主调度智能体 - 系统的协调中心
负责任务分解、智能体调度和结果整合
"""

import asyncio
import json
from typing import List, Dict, Any, Optional
from datetime import datetime
import uuid

from ...core.interfaces import (
    OrchestratorInterface, AgentInterface, ToolInterface, 
    Task, TaskStatus
)
from ...core.events import EventEmitter, SystemEvents
from ...core.container import container
from ...lightrag_engine.rag_controller import RAGController


class MainOrchestrator(OrchestratorInterface, EventEmitter):
    """主调度智能体 - 任务协调和分发"""
    
    def __init__(self):
        super().__init__()
        self.registered_agents: Dict[str, AgentInterface] = {}
        self.registered_tools: Dict[str, ToolInterface] = {}
        self.rag_controller: Optional[RAGController] = None
        self.task_queue: List[Task] = []
        self.running_tasks: Dict[str, Task] = {}
        self._initialized = False
    
    async def initialize(self):
        """初始化调度器"""
        if self._initialized:
            return
        
        try:
            # 获取RAG控制器
            self.rag_controller = container.rag_controller()
            # Mock RAG控制器不需要初始化
            if hasattr(self.rag_controller, 'initialize'):
                await self.rag_controller.initialize()
            
            self._initialized = True
            
            await self.emit("orchestrator_initialized", {
                "registered_agents": len(self.registered_agents),
                "registered_tools": len(self.registered_tools)
            })
            
        except Exception as e:
            await self.emit_error("orchestrator_initialization", e)
            raise
    
    def register_agent(self, agent: AgentInterface):
        """注册智能体"""
        self.registered_agents[agent.name] = agent
    
    def register_tool(self, tool: ToolInterface):
        """注册工具"""
        self.registered_tools[tool.name] = tool
    
    async def orchestrate_task(self, task: Task) -> Task:
        """协调任务执行"""
        if not self._initialized:
            await self.initialize()
        
        try:
            await self.emit(SystemEvents.TASK_RECEIVED, {
                "task_id": task.id,
                "task_type": task.task_type,
                "description": task.description[:100] + "..." if len(task.description) > 100 else task.description
            })
            
            # 更新任务状态
            task.status = TaskStatus.IN_PROGRESS
            task.metadata = task.metadata or {}
            task.metadata["started_at"] = datetime.now().isoformat()
            
            self.running_tasks[task.id] = task
            
            await self.emit(SystemEvents.TASK_STARTED, {
                "task_id": task.id,
                "task_type": task.task_type
            })
            
            # 任务分析和分解
            execution_plan = await self._analyze_and_plan_task(task)
            
            # 执行任务计划
            result = await self._execute_task_plan(task, execution_plan)
            
            # 更新任务结果
            task.result = result
            task.status = TaskStatus.COMPLETED
            task.metadata["completed_at"] = datetime.now().isoformat()
            task.metadata["execution_plan"] = execution_plan
            
            # 清理运行中的任务
            if task.id in self.running_tasks:
                del self.running_tasks[task.id]
            
            await self.emit(SystemEvents.TASK_COMPLETED, {
                "task_id": task.id,
                "task_type": task.task_type,
                "execution_time": self._calculate_execution_time(task)
            })
            
            return task
            
        except Exception as e:
            task.status = TaskStatus.FAILED
            task.result = {"error": str(e)}
            task.metadata["failed_at"] = datetime.now().isoformat()
            task.metadata["error"] = str(e)
            
            if task.id in self.running_tasks:
                del self.running_tasks[task.id]
            
            await self.emit(SystemEvents.TASK_FAILED, {
                "task_id": task.id,
                "task_type": task.task_type,
                "error": str(e)
            })
            
            return task
    
    async def _analyze_and_plan_task(self, task: Task) -> Dict[str, Any]:
        """分析任务并制定执行计划"""
        try:
            # 使用RAG分析任务类型和要求
            rag_analysis = await self.rag_controller.process_query(
                query=f"任务分析：{task.description}",
                template_name="general_qa",
                top_k=5,
                return_intermediate=False
            )
            
            # 基于任务描述确定执行策略
            execution_plan = {
                "task_id": task.id,
                "analysis": rag_analysis.get("answer", ""),
                "strategy": self._determine_execution_strategy(task),
                "required_agents": [],
                "required_tools": [],
                "steps": [],
                "estimated_time": 0
            }
            
            # 根据任务类型制定具体计划
            if task.task_type == "resume_analysis":
                execution_plan.update(await self._plan_resume_analysis_task(task))
            elif task.task_type == "candidate_comparison":
                execution_plan.update(await self._plan_candidate_comparison_task(task))
            elif task.task_type == "skill_assessment":
                execution_plan.update(await self._plan_skill_assessment_task(task))
            elif task.task_type == "background_check":
                execution_plan.update(await self._plan_background_check_task(task))
            else:
                execution_plan.update(await self._plan_general_task(task))
            
            return execution_plan
            
        except Exception as e:
            await self.emit_error("analyze_and_plan_task", e)
            return {
                "task_id": task.id,
                "strategy": "fallback",
                "required_agents": [],
                "steps": [{"action": "direct_rag_query", "description": task.description}],
                "error": str(e)
            }
    
    def _determine_execution_strategy(self, task: Task) -> str:
        """确定执行策略"""
        task_keywords = task.description.lower()
        
        if any(keyword in task_keywords for keyword in ["简历", "resume", "候选人", "candidate"]):
            return "resume_focused"
        elif any(keyword in task_keywords for keyword in ["比较", "对比", "compare"]):
            return "comparative_analysis"
        elif any(keyword in task_keywords for keyword in ["技能", "skill", "能力", "ability"]):
            return "skill_analysis"
        elif any(keyword in task_keywords for keyword in ["背景", "验证", "background", "verify"]):
            return "background_verification"
        else:
            return "general_query"
    
    async def _plan_resume_analysis_task(self, task: Task) -> Dict[str, Any]:
        """规划简历分析任务"""
        return {
            "required_agents": ["resume_analyst", "kg_analyst"],
            "required_tools": ["rag_search", "entity_extraction"],
            "steps": [
                {
                    "step": 1,
                    "agent": "resume_analyst",
                    "action": "analyze_resume",
                    "description": "分析简历基本信息和结构",
                    "tools": ["rag_search"]
                },
                {
                    "step": 2,
                    "agent": "kg_analyst", 
                    "action": "extract_entities_and_relations",
                    "description": "提取实体和关系信息",
                    "tools": ["entity_extraction"]
                },
                {
                    "step": 3,
                    "agent": "resume_analyst",
                    "action": "generate_comprehensive_report",
                    "description": "生成综合分析报告",
                    "tools": ["rag_search"]
                }
            ],
            "estimated_time": 30
        }
    
    async def _plan_candidate_comparison_task(self, task: Task) -> Dict[str, Any]:
        """规划候选人比较任务"""
        return {
            "required_agents": ["resume_analyst", "potential_forecaster"],
            "required_tools": ["rag_search", "comparison_analysis"],
            "steps": [
                {
                    "step": 1,
                    "agent": "resume_analyst",
                    "action": "analyze_multiple_candidates",
                    "description": "分析多个候选人简历",
                    "tools": ["rag_search"]
                },
                {
                    "step": 2,
                    "agent": "potential_forecaster",
                    "action": "assess_potential",
                    "description": "评估候选人潜力",
                    "tools": ["potential_assessment"]
                },
                {
                    "step": 3,
                    "agent": "resume_analyst",
                    "action": "generate_comparison_report",
                    "description": "生成比较分析报告",
                    "tools": ["comparison_analysis"]
                }
            ],
            "estimated_time": 45
        }
    
    async def _plan_skill_assessment_task(self, task: Task) -> Dict[str, Any]:
        """规划技能评估任务"""
        return {
            "required_agents": ["resume_analyst", "kg_analyst"],
            "required_tools": ["rag_search", "skill_matching"],
            "steps": [
                {
                    "step": 1,
                    "agent": "kg_analyst",
                    "action": "extract_skills",
                    "description": "提取技能信息",
                    "tools": ["entity_extraction"]
                },
                {
                    "step": 2,
                    "agent": "resume_analyst",
                    "action": "assess_skill_level",
                    "description": "评估技能水平",
                    "tools": ["skill_matching"]
                }
            ],
            "estimated_time": 20
        }
    
    async def _plan_background_check_task(self, task: Task) -> Dict[str, Any]:
        """规划背景核查任务"""
        return {
            "required_agents": ["background_verifier"],
            "required_tools": ["web_search", "linkedin_api"],
            "steps": [
                {
                    "step": 1,
                    "agent": "background_verifier",
                    "action": "verify_background",
                    "description": "验证背景信息",
                    "tools": ["web_search", "linkedin_api"]
                }
            ],
            "estimated_time": 60
        }
    
    async def _plan_general_task(self, task: Task) -> Dict[str, Any]:
        """规划通用任务"""
        return {
            "required_agents": [],
            "required_tools": ["rag_search"],
            "steps": [
                {
                    "step": 1,
                    "action": "direct_rag_query",
                    "description": "直接使用RAG查询",
                    "tools": ["rag_search"]
                }
            ],
            "estimated_time": 15
        }
    
    async def _execute_task_plan(self, task: Task, execution_plan: Dict[str, Any]) -> Dict[str, Any]:
        """执行任务计划"""
        try:
            results = {
                "plan_id": execution_plan.get("task_id"),
                "strategy": execution_plan.get("strategy"),
                "step_results": [],
                "final_result": None,
                "execution_summary": {}
            }
            
            steps = execution_plan.get("steps", [])
            
            for step in steps:
                step_result = await self._execute_step(task, step)
                results["step_results"].append({
                    "step": step.get("step", len(results["step_results"]) + 1),
                    "action": step.get("action"),
                    "result": step_result,
                    "timestamp": datetime.now().isoformat()
                })
            
            # 整合所有步骤结果
            results["final_result"] = await self._integrate_step_results(task, results["step_results"])
            
            # 生成执行摘要
            results["execution_summary"] = {
                "total_steps": len(steps),
                "successful_steps": sum(1 for r in results["step_results"] if not r["result"].get("error")),
                "agents_used": list(set(step.get("agent") for step in steps if step.get("agent"))),
                "tools_used": list(set(tool for step in steps for tool in step.get("tools", []))),
                "total_execution_time": self._calculate_execution_time(task)
            }
            
            return results
            
        except Exception as e:
            await self.emit_error("execute_task_plan", e)
            return {
                "error": str(e),
                "plan_id": execution_plan.get("task_id"),
                "partial_results": results.get("step_results", [])
            }
    
    async def _execute_step(self, task: Task, step: Dict[str, Any]) -> Dict[str, Any]:
        """执行单个步骤"""
        try:
            action = step.get("action")
            agent_name = step.get("agent")
            tools = step.get("tools", [])
            
            if action == "direct_rag_query":
                # 直接使用RAG查询
                return await self._execute_rag_query(task.description)
            
            elif agent_name and agent_name in self.registered_agents:
                # 使用指定智能体
                agent = self.registered_agents[agent_name]
                
                # 创建子任务
                sub_task = Task(
                    id=f"{task.id}_step_{step.get('step', 0)}",
                    description=step.get("description", task.description),
                    task_type=action,
                    metadata={
                        "parent_task_id": task.id,
                        "step_info": step,
                        "tools": tools
                    }
                )
                
                # 执行子任务
                result_task = await agent.process_task(sub_task)
                return {
                    "success": result_task.status == TaskStatus.COMPLETED,
                    "result": result_task.result,
                    "agent": agent_name,
                    "task_id": sub_task.id
                }
            
            else:
                # 降级处理 - 使用RAG查询
                return await self._execute_rag_query(task.description)
                
        except Exception as e:
            await self.emit_error("execute_step", e)
            return {
                "success": False,
                "error": str(e),
                "step_info": step
            }
    
    async def _execute_rag_query(self, query: str) -> Dict[str, Any]:
        """执行RAG查询"""
        try:
            result = await self.rag_controller.process_query(
                query=query,
                return_intermediate=True
            )
            
            return {
                "success": True,
                "result": result,
                "method": "rag_query"
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "method": "rag_query"
            }
    
    async def _integrate_step_results(self, task: Task, step_results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """整合步骤结果"""
        try:
            # 收集所有成功的结果
            successful_results = [r for r in step_results if r["result"].get("success", False)]
            
            if not successful_results:
                return {"error": "No successful steps to integrate"}
            
            # 构建整合结果
            integrated_result = {
                "task_id": task.id,
                "task_description": task.description,
                "integration_method": "sequential_synthesis",
                "component_results": [],
                "final_answer": ""
            }
            
            # 收集组件结果
            all_answers = []
            for step_result in successful_results:
                result_data = step_result["result"]["result"]
                
                if isinstance(result_data, dict) and "answer" in result_data:
                    all_answers.append(result_data["answer"])
                    integrated_result["component_results"].append({
                        "step": step_result["step"],
                        "action": step_result["action"],
                        "answer": result_data["answer"],
                        "metadata": result_data.get("metadata", {})
                    })
                elif isinstance(result_data, str):
                    all_answers.append(result_data)
                    integrated_result["component_results"].append({
                        "step": step_result["step"],
                        "action": step_result["action"],
                        "answer": result_data
                    })
            
            # 生成最终答案
            if len(all_answers) == 1:
                integrated_result["final_answer"] = all_answers[0]
            else:
                # 使用RAG生成整合答案
                synthesis_query = f"基于以下分析结果，请提供综合答案：\n\n" + "\n\n".join(all_answers)
                synthesis_result = await self.rag_controller.process_query(
                    query=synthesis_query,
                    template_name="general_qa"
                )
                integrated_result["final_answer"] = synthesis_result.get("answer", "整合结果时出现错误")
            
            return integrated_result
            
        except Exception as e:
            await self.emit_error("integrate_step_results", e)
            return {
                "error": str(e),
                "partial_integration": step_results
            }
    
    def _calculate_execution_time(self, task: Task) -> float:
        """计算执行时间"""
        try:
            if not task.metadata:
                return 0.0
            
            started_at = task.metadata.get("started_at")
            completed_at = task.metadata.get("completed_at") or task.metadata.get("failed_at")
            
            if started_at and completed_at:
                start_time = datetime.fromisoformat(started_at)
                end_time = datetime.fromisoformat(completed_at)
                return (end_time - start_time).total_seconds()
            
            return 0.0
            
        except Exception:
            return 0.0
    
    async def get_orchestrator_status(self) -> Dict[str, Any]:
        """获取调度器状态"""
        return {
            "initialized": self._initialized,
            "registered_agents": list(self.registered_agents.keys()),
            "registered_tools": list(self.registered_tools.keys()),
            "task_queue_size": len(self.task_queue),
            "running_tasks": len(self.running_tasks),
            "running_task_ids": list(self.running_tasks.keys())
        }
    
    async def cancel_task(self, task_id: str) -> bool:
        """取消任务"""
        try:
            if task_id in self.running_tasks:
                task = self.running_tasks[task_id]
                task.status = TaskStatus.FAILED
                task.result = {"error": "Task cancelled by user"}
                task.metadata["cancelled_at"] = datetime.now().isoformat()
                
                del self.running_tasks[task_id]
                
                await self.emit("task_cancelled", {
                    "task_id": task_id
                })
                
                return True
            
            return False
            
        except Exception as e:
            await self.emit_error("cancel_task", e)
            return False